Skip to content

Conversation

@gabotechs
Copy link
Collaborator

This PR provides a better organization for the project colocating into modules the different pieces based on the features they are providing.

The new folder structure, as it colocates things related to each other more efficiently, allows to tweak the privacy of some structs/methods so that they are not expose to outside unrelated modules, and therefore reducing the chances of future coupling.

All the code is mostly untouched, only a couple of code changes are shipped for making this possible:

  • The ExecutionTask now has two forms:
    • ExecutionTask: struct with the Url already parsed and the partition groups already as usizes, so that we don't need to do juggling in the code for parsing Urls or casing u64s to usizes
    • ExecutionTaskProto: the protobuf representation, decoupled from the actual useful struct
  • The test in the do_get method was using directly the proto representation of the ExecutionStage and the ExecutionTask, but they are no longer exposed to the outside, so instead the normal structs are used, which then are converted to protos with the proto_to_stage function. This actually results in a nice cleanup that saves some lines of code.
common/
    mod.rs
    composed_extension_codec.rs
    ttl_map.rs
errors/
    mod.rs
    arrow_error.rs
    datafusion_error.rs
    io_error.rs
    objectstore_error.rs
    parquet_error.rs
    parser_error.rs
    schema_error.rs
execution_plans/
    mod.rs
    arrow_flight_read.rs
    partition_isolator.rs
    stage.rs
flight_service/
    mod.rs
    do_get.rs
    service.rs
    session_builder.rs
protobuf/
    mod.rs
    distributed_codec.rs
    stage_proto.rs
    user_codec.rs
test_utils/
    mod.rs
    insta.rs
    localhost.rs
    mock_exec.rs
    parquet.rs
    tpch.rs
lib.rs
channel_resolver_ext.rs
config_extension_ext.rs
distributed_ext.rs
distributed_physical_optimizer_rule.rs

@gabotechs gabotechs force-pushed the gabrielmusat/bigger-tpch-tests branch 2 times, most recently from 898fbf0 to c01c11c Compare September 5, 2025 06:03
Base automatically changed from gabrielmusat/bigger-tpch-tests to main September 5, 2025 06:10
@gabotechs gabotechs force-pushed the gabrielmusat/cleanup branch from 8d7ba36 to 355da32 Compare September 5, 2025 06:11
Copy link
Collaborator

@NGA-TRAN NGA-TRAN left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice refactor and rename.

mod tests {
use crate::assert_snapshot;
use crate::physical_optimizer::DistributedPhysicalOptimizerRule;
use crate::distributed_physical_optimizer_rule::DistributedPhysicalOptimizerRule;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like you rename it to distributed_physical_optimizer_rule. Make it clearer

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, I've tried to rename files to be a bit more consistent with the content.

) -> Result<Arc<dyn ExecutionPlan>> {
// We can only optimize plans that are not already distributed
if plan.as_any().is::<ExecutionStage>() {
if plan.as_any().is::<StageExec>() {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

StageExec makes it more like DF style 👍

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, it seems like an unwritten rule to suffix all ExecutionPlan implementations with *Exec

// Helper to create a mock physical plan
fn create_mock_physical_plan(partitions: usize) -> Arc<dyn ExecutionPlan> {
let node = Arc::new(EmptyExec::new(SchemaRef::new(Schema::empty())));
Arc::new(RepartitionExec::try_new(node, Partitioning::RoundRobinBatch(partitions)).unwrap())
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn’t directly related to this PR, but might be something we consider updating or adding in a follow-up.

This RoundRobinBatch got me thinking 🤔. On one hand, it's a valid scenario we want our tests to cover. On the other, round-robin repartitioning isn’t ideal from a performance standpoint. How about we introduce a variety of partitioning schemes in our mock tests? That way, we can surface more realistic and appropriate cases while still including edge scenarios like this one.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, we should do something with RoundRobin repartitions...

We probably should extend our parquet files test datasets so that we can cover more cases and more esoteric distribution patterns.

For example, it would be nice to have two test parquet files with related data that we can use to stress joins in weird ways.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably something to follow up soon

@gabotechs gabotechs merged commit b728287 into main Sep 5, 2025
4 checks passed
@gabotechs gabotechs deleted the gabrielmusat/cleanup branch September 5, 2025 15:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants